import numpy as np
import matplotlib

matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
import pandas as pd
from math import isclose
import random
from collections import Counter
import numpy.random as npr

np.set_printoptions(precision=4, suppress=True)


def simulate_chain(P, start, n_steps):
    """
    Simulate Markov chain with transition matrix P starting from state index start
    for n_steps.
    Returns the list of visted states including the start state.
    :param P:
    :param start:
    :param n_steps:
    :return:
    """
    states = [start]
    cur = start
    for _ in range(n_steps):
        cur = npr.choice(a=len(P), p=P[cur])
        states.append(cur)

    return states


print("#" * 80)

print("Section 0:Memoryless property")
print("- We simulate two different histories that lead to the same current state and show"
      " future distribution depends only on current state")
P0 = np.array([
    [0.7, 0.3],
    [0.2, 0.8]
])
print("Transition matrix P0:\n", P0)


# simulate many trajectories and filter those hitting state 1 at time 2 from different
# starts/history
def future_distribution_cond_on_current(P, current_state, trials=20000, future_steps=1):
    counts = Counter()
    for _ in range(trials):
        # start uniformly at random;we only care if at time 2 we are in current_state
        # 只保留在时间2达到状态1的路径，这些路径可能有不同的历史：0→0→1、0→1→1、1→0→1、1→1→1等
        path = simulate_chain(P, npr.choice(len(P)), 2)
        if path[-1] == current_state:
            # 对所有达到状态1的路径，观察下一步的状态
            # 计算条件概率：P(X₃=s | X₂=1)
            future = simulate_chain(P, current_state, future_steps)[1]  # one-step ahead
            counts[future] += 1
    total = sum(counts.values())
    for s in range(len(P)):
        # P(X_{t+1} = j | X_t = i, X_{t-1} = i_{t-1}, ..., X_0 = i_0) = P(X_{t+1} = j | X_t = i)
        print(f"P(X_{{t+1}}={s} | X_t={current_state}) ≈ {counts[s] / total:.4f}")

print("Estimate conditional future distribution given current state 1:")
future_distribution_cond_on_current(P0, 1)

